package transform

import (
	"context"
	"crypto/rand"
	"errors"
	"io"
	mathRand "math/rand/v2"
	"net"
	"net/http"
	"smart-flow/agent/proto"
	"sync/atomic"
	"time"
)

type QuitFunc func()

type DataSender struct {
	rateDelta       int64        // 速率浮动系数
	tcpPerUseCount  int          // 每次使用多少个地址
	udpPerUseCount  int          // 每次使用多少个地址
	httpPerUseCount int          // 每次使用多少个地址
	tcpAddresses    []string     // Tcp地址列表
	udpAddresses    []string     // Udp地址列表
	httpUrls        []string     // 目标Http地址列表
	dataPool        []byte       // 只读数据池
	poolSize        int64        // 数据池大小
	blockCount      int64        // 每次发送的数据倍率(blockCount*poolSize)
	tcpRateLimit    atomic.Int64 // 每秒限制速率(Bps)
	tcpBaseRate     atomic.Int64 // 每秒基准速率(限制速率/perUseCount)
	tcpRunRate      atomic.Int64 // 每秒运行速率
	udpRateLimit    atomic.Int64 // 每秒限制速率(Bps)
	udpBaseRate     atomic.Int64 // 每秒基准速率(限制速率/perUseCount)
	udpRunRate      atomic.Int64 // 每秒运行速率
	httpRateLimit   atomic.Int64 // 每秒限制速率(Bps)
	httpBaseRate    atomic.Int64 // 每秒基准速率(限制速率/perUseCount)
	httpRunRate     atomic.Int64 // 每秒运行速率
	paused          atomic.Bool  // 暂停
	stoped          atomic.Bool  // 停止
	httpClient      *http.Client
	quitFunc        QuitFunc
}

func (ds *DataSender) Init(data *proto.LoginResp, poolSize int64, quitFunc QuitFunc) *DataSender {
	dataPool := make([]byte, poolSize) // 生成随机数据池
	if _, err := rand.Read(dataPool); err != nil {
		return nil
	}
	ds.rateDelta = data.RateDelta
	ds.tcpAddresses = data.TcpAddresses
	ds.udpAddresses = data.UdpAddresses
	ds.httpUrls = data.HttpUrls
	ds.dataPool = dataPool
	ds.poolSize = poolSize
	ds.blockCount = data.BlockCount
	ds.quitFunc = quitFunc
	ds.UpdatePerUseCount(data.PerUseCount)
	ds.UpdateRates(data.TcpRateLimit, data.UdpRateLimit, data.HttpRateLimit)
	if len(data.HttpUrls) > 0 {
		ds.httpClient = &http.Client{
			Transport: &http.Transport{
				DialContext: (&net.Dialer{
					Timeout: 30 * time.Second, // 设置连接超时时间
				}).DialContext,
				MaxIdleConns:        50,
				MaxIdleConnsPerHost: 50,
				IdleConnTimeout:     180 * time.Second,
				WriteBufferSize:     32 * 1024,
			},
			Timeout: 0, // 禁用整体超时
		}
	}
	return ds
}

func (ds *DataSender) UpdateRates(tcpRateLimit, udpRateLimit, httpRateLimit int64) {
	if len(ds.tcpAddresses) > 0 && tcpRateLimit > 0 && ds.tcpPerUseCount > 0 {
		tcpRateLimit = tcpRateLimit * KBit / Byte2Bit
		ds.tcpRateLimit.Store(tcpRateLimit)
		ds.tcpBaseRate.Store(tcpRateLimit / int64(ds.tcpPerUseCount))
	}
	if len(ds.udpAddresses) > 0 && udpRateLimit > 0 && ds.udpPerUseCount > 0 {
		udpRateLimit = udpRateLimit * KBit / Byte2Bit
		ds.udpRateLimit.Store(udpRateLimit)
		ds.udpBaseRate.Store(udpRateLimit / int64(ds.udpPerUseCount))
	}
	if len(ds.httpUrls) > 0 && httpRateLimit > 0 && ds.httpPerUseCount > 0 {
		httpRateLimit = httpRateLimit * KBit / Byte2Bit
		ds.httpRateLimit.Store(httpRateLimit)
		ds.httpBaseRate.Store(httpRateLimit / int64(ds.httpPerUseCount))
	}
}

func (ds *DataSender) UpdatePerUseCount(perUseCount int) {
	ds.tcpPerUseCount = max(perUseCount, min(perUseCount, len(ds.tcpAddresses)))
	ds.udpPerUseCount = max(perUseCount, min(perUseCount, len(ds.udpAddresses)))
	ds.httpPerUseCount = max(perUseCount, min(perUseCount, len(ds.httpUrls)))
}

func (ds *DataSender) UpdateRateDelta(reteDelta int64) {
	ds.rateDelta = reteDelta
}

func (ds *DataSender) Start() {
	length := len(ds.tcpAddresses)
	if length > 0 && ds.tcpRateLimit.Load() > 0 {
		go ds.startTcp(length)
	}
	length = len(ds.udpAddresses)
	if length > 0 && ds.udpRateLimit.Load() > 0 {
		go ds.startUdp(length)
	}
	length = len(ds.httpUrls)
	if length > 0 && ds.httpRateLimit.Load() > 0 {
		go ds.startHttp(length)
	}
}

func (ds *DataSender) startTcp(length int) {
	i := 0
	availableRate, rate, baseRate := int64(0), int64(0), int64(0)
	rateLimit := ds.tcpRateLimit.Load()
	for {
		baseRate = ds.tcpBaseRate.Load()
		availableRate = rateLimit - ds.tcpRunRate.Load()
		if availableRate < baseRate/2 { // 可用速率小于base/2不会运行, 应对随机总速率后的微增, 确保正常分片完成后再运行
			time.Sleep(time.Second)
			continue
		}
		rate = randomShared(baseRate, availableRate, ds.rateDelta)
		// fmt.Printf("tcp 运行速率: %d Kbps, 速率限制: %d Kbps, 可用速率: %d Kbps\n", rate*Byte2Bit/KBit, rateLimit*Byte2Bit/KBit, availableRate*Byte2Bit/KBit)
		ds.tcpRunRate.Add(rate)
		go ds.sendData2Tcp(ds.tcpAddresses[i], rate)
		i++
		if i%ds.tcpPerUseCount == 0 {
			rateLimit = ds.tcpRateLimit.Load()
			rateLimit = randomRate(rateLimit, ds.rateDelta)
		}
		if i == length {
			i = 0
		}
	}
}

func (ds *DataSender) sendData2Tcp(address string, rate int64) {
	defer ds.tcpRunRate.Add(-rate)
	addr, err := net.ResolveTCPAddr("tcp", address)
	if err != nil {
		ds.quitFunc()
		return
	}
	conn, err := net.DialTCP("tcp", nil, addr)
	if err != nil {
		ds.quitFunc()
		return
	}
	defer func(conn *net.TCPConn) {
		_ = conn.Close()
	}(conn)
	totalSize := ds.poolSize * ds.blockCount / int64(ds.tcpPerUseCount) // 保证每次在不长的时间内完成(30s~3m30s)
	// fmt.Printf("tcp 任务内共需发送: %d KB, 端口:%d, 速率限制: %d Kbps\n", totalSize/KB, port, rate*Byte2Bit/KBit)
	totalSent := int64(0)
	bytesInSec := int64(0)
	lastTime := time.Now()
	for totalSent < totalSize {
		if ds.Stoped() {
			return
		}
		if ds.Paused() {
			time.Sleep(time.Second)
			continue
		}
		now := time.Now()
		elapsed := now.Sub(lastTime)
		if elapsed >= time.Second {
			lastTime = now
			// fmt.Printf("tcp 已发送 %d KB，端口:%d, 速率: %d Kbps, 速率限制: %d Kbps\n", totalSent/KB, port, bytesInSec*Byte2Bit/KBit, rate*Byte2Bit/KBit)
			bytesInSec = 0 // 每秒更新一次计数器
		}
		toWrite := rate - bytesInSec // 计算本次发送的数据量
		if toWrite <= 0 {
			time.Sleep(time.Second - elapsed) // 如果当前秒已经用完配额，则等待下一轮
			continue
		}
		packetSize := min(500, toWrite) // 确定单次发送的数据大小
		startPos := totalSent % ds.poolSize
		endPos := startPos + packetSize
		if endPos > ds.poolSize {
			startPos = 0
			endPos = packetSize
		}
		_, err = conn.Write(ds.dataPool[startPos:endPos]) // 发送UDP数据
		if err != nil {
			ds.quitFunc()
			return
		}
		totalSent += packetSize
		bytesInSec += packetSize
	}
}

func (ds *DataSender) startUdp(length int) {
	i := 0
	availableRate, rate, baseRate := int64(0), int64(0), int64(0)
	rateLimit := ds.udpRateLimit.Load()
	for {
		baseRate = ds.udpBaseRate.Load()
		availableRate = rateLimit - ds.udpRunRate.Load()
		if availableRate < baseRate/2 { // 可用速率小于base/2不会运行, 应对随机总速率后的微增, 确保正常分片完成后再运行
			time.Sleep(time.Second)
			continue
		}
		rate = randomShared(baseRate, availableRate, ds.rateDelta)
		// fmt.Printf("udp 运行速率: %d Kbps, 速率限制: %d Kbps, 可用速率: %d Kbps\n", rate*Byte2Bit/KBit, rateLimit*Byte2Bit/KBit, availableRate*Byte2Bit/KBit)
		ds.udpRunRate.Add(rate)
		go ds.sendData2Udp(ds.udpAddresses[i], rate)
		i++
		if i%ds.udpPerUseCount == 0 {
			rateLimit = ds.udpRateLimit.Load()
			rateLimit = randomRate(rateLimit, ds.rateDelta)
		}
		if i == length {
			i = 0
		}
	}
}

func (ds *DataSender) sendData2Udp(address string, rate int64) {
	defer ds.udpRunRate.Add(-rate)
	addr, err := net.ResolveUDPAddr("udp", address)
	if err != nil {
		ds.quitFunc()
		return
	}
	conn, err := net.DialUDP("udp", nil, addr)
	if err != nil {
		ds.quitFunc()
		return
	}
	defer func(conn *net.UDPConn) {
		_ = conn.Close()
	}(conn)
	totalSize := ds.poolSize * ds.blockCount / int64(ds.udpPerUseCount) // 保证每次在不长的时间内完成(30s~3m30s)
	// fmt.Printf("udp 任务内共需发送: %d KB, 端口:%d, 速率限制: %d Kbps\n", totalSize/KB, port, rate*Byte2Bit/KBit)
	totalSent := int64(0)
	bytesInSec := int64(0)
	lastTime := time.Now()
	for totalSent < totalSize {
		if ds.Stoped() {
			return
		}
		if ds.Paused() {
			time.Sleep(time.Second)
			continue
		}
		now := time.Now()
		elapsed := now.Sub(lastTime)
		if elapsed >= time.Second {
			lastTime = now
			// fmt.Printf("udp 已发送 %d KB，端口:%d, 速率: %d Kbps, 速率限制: %d Kbps\n", totalSent/KB, port, bytesInSec*Byte2Bit/KBit, rate*Byte2Bit/KBit)
			bytesInSec = 0 // 每秒更新一次计数器
		}
		toWrite := rate - bytesInSec // 计算本次发送的数据量
		if toWrite <= 0 {
			time.Sleep(time.Second - elapsed) // 如果当前秒已经用完配额，则等待下一轮
			continue
		}
		packetSize := min(500, toWrite) // 确定单次发送的数据大小
		startPos := totalSent % ds.poolSize
		endPos := startPos + packetSize
		if endPos > ds.poolSize {
			startPos = 0
			endPos = packetSize
		}
		_, err = conn.Write(ds.dataPool[startPos:endPos]) // 发送UDP数据
		if err != nil {
			ds.quitFunc()
			return
		}
		totalSent += packetSize
		bytesInSec += packetSize
	}
}

func (ds *DataSender) startHttp(length int) {
	i := 0
	availableRate, rate, baseRate := int64(0), int64(0), int64(0)
	rateLimit := ds.httpRateLimit.Load()
	for {
		baseRate = ds.httpBaseRate.Load()
		availableRate = rateLimit - ds.httpRunRate.Load()
		if availableRate < baseRate/2 { // 可用速率小于base/2不会运行, 应对随机总速率后的微增, 确保正常分片完成后再运行
			time.Sleep(time.Second)
			continue
		}
		rate = randomShared(baseRate, availableRate, ds.rateDelta)
		// fmt.Printf("http 运行速率: %d Kbps, 速率限制: %d Kbps, 可用速率: %d Kbps\n", rate*Byte2Bit/KBit, rateLimit*Byte2Bit/KBit, availableRate*Byte2Bit/KBit)
		ds.httpRunRate.Add(rate)
		go ds.sendData2Http(ds.httpUrls[i], rate)
		i++
		if i%ds.httpPerUseCount == 0 {
			rateLimit = ds.httpRateLimit.Load()
			rateLimit = randomRate(rateLimit, ds.rateDelta)
		}
		if i == length {
			i = 0
		}
	}
}

func (ds *DataSender) sendData2Http(url string, rate int64) {
	defer ds.httpRunRate.Add(-rate)
	pr, pw := io.Pipe()
	defer func(pr *io.PipeReader) {
		_ = pr.Close()
	}(pr)
	go func() {
		ds.sendDataWithPipe(pw, rate)
	}()
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	req, err := http.NewRequestWithContext(ctx, "POST", url, pr)
	if err != nil {
		ds.quitFunc()
		return
	}
	req.Header.Set("Content-Type", "application/octet-stream") // 标准二进制流类型
	req.Header.Set("Content-Transfer-Encoding", "binary")
	resp, err := ds.httpClient.Do(req)
	if err != nil {
		ds.quitFunc()
		return
	}
	defer func(resp *http.Response) {
		if resp != nil && resp.Body != nil {
			_ = resp.Body.Close()
		}
	}(resp)
	_, _ = io.Copy(io.Discard, resp.Body) // 确保消费完整个 body，避免 pipe 提前关闭
}

func (ds *DataSender) sendDataWithPipe(pw *io.PipeWriter, rate int64) {
	defer func(pw *io.PipeWriter) {
		_ = pw.Close()
	}(pw)
	totalSize := ds.poolSize * ds.blockCount / int64(ds.httpPerUseCount) // 保证每次在不长的时间内完成(30s~3m30s)
	// fmt.Printf("http 任务内共需发送: %d KB, 速率限制: %d Kbps\n", totalSize/KB, rate*Byte2Bit/KBit)
	totalSent := int64(0)
	bytesInSec := int64(0)
	lastTime := time.Now()
	for totalSent < totalSize {
		if ds.Stoped() {
			return
		}
		if ds.Paused() {
			time.Sleep(time.Second)
			continue
		}
		now := time.Now()
		elapsed := now.Sub(lastTime)
		if elapsed >= time.Second {
			lastTime = now
			// fmt.Printf("http 已发送 %d KB，速率: %d Kbps, 速率限制: %d Kbps\n", totalSent/KB, bytesInSec*Byte2Bit/KBit, rate*Byte2Bit/KBit)
			bytesInSec = 0 // 每秒更新一次计数器
		}
		toWrite := rate - bytesInSec // 计算本次发送的数据量
		if toWrite <= 0 {
			time.Sleep(time.Second - elapsed) // 如果当前秒已经用完配额，则等待下一轮
			continue
		}
		packetSize := min(ds.poolSize, toWrite) // 确定单次发送的数据大小
		startPos := totalSent % ds.poolSize
		endPos := startPos + packetSize
		if endPos > ds.poolSize {
			startPos = 0
			endPos = packetSize
		}
		n, err := pw.Write(ds.dataPool[startPos:endPos])
		if err != nil {
			if errors.Is(err, io.ErrClosedPipe) { // 检查是否是因为pipe关闭导致的错误
				return
			}
			ds.quitFunc()
			return
		}
		totalSent += int64(n)
		bytesInSec += int64(n)
	}
}

func (ds *DataSender) Pause() {
	ds.paused.Store(true)
}

func (ds *DataSender) Resume() {
	ds.paused.Store(false)
}

func (ds *DataSender) Paused() bool {
	return ds.paused.Load()
}

func (ds *DataSender) Stop() {
	ds.stoped.Store(true)
}

func (ds *DataSender) Stoped() bool {
	return ds.stoped.Load()
}

func randomShared(base, availableRate, rateDelta int64) int64 {
	var rate int64
	maxRate := base * (100 + rateDelta) / 100
	if availableRate <= maxRate+base/2 { // 确保不会出现很小的速率, 因为小于base/2不会执行
		rate = availableRate
	} else {
		deltaPercent := mathRand.Int64N(rateDelta + 1)
		rate = min(availableRate, base*(100+deltaPercent)/100) // base~base(100+rateDelta)%
	}
	return rate
}

func randomRate(rateLimit, rateDelta int64) int64 {
	if rateDelta <= 0 {
		return rateLimit
	}
	deltaPercent := mathRand.Int64N(rateDelta*2+1) - rateDelta
	return rateLimit * (100 + deltaPercent) / 100
}
